# Copyright (c) OpenMMLab. All rights reserved.
import argparse

import torch
from peft import PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    CLIPImageProcessor,
    CLIPVisionModel,
)

from xtuner.model.utils import LoadWoInit
from xtuner.utils.device import get_device_name


def parse_args():
    parser = argparse.ArgumentParser(
        description="Merge a HuggingFace adapter to base model"
    )
    parser.add_argument("model_name_or_path", help="model name or path")
    parser.add_argument("adapter_name_or_path", help="adapter name or path")
    parser.add_argument("save_dir", help="the directory to save the merged model")
    parser.add_argument(
        "--max-shard-size",
        type=str,
        default="2GB",
        help="Only applicable for LLM. The maximum size for "
        "each sharded checkpoint.",
    )
    parser.add_argument(
        "--is-clip", action="store_true", help="Indicate if the model is a clip model"
    )
    parser.add_argument(
        "--safe-serialization",
        action="store_true",
        help="Indicate if using `safe_serialization`",
    )
    parser.add_argument(
        "--device",
        default=get_device_name(),
        choices=("cuda", "cpu", "npu", "auto"),
        help="Indicate the device",
    )

    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    if args.is_clip:
        with LoadWoInit():
            model = CLIPVisionModel.from_pretrained(
                args.model_name_or_path, device_map=args.device
            )
        processor = CLIPImageProcessor.from_pretrained(args.model_name_or_path)
    else:
        with LoadWoInit():
            model = AutoModelForCausalLM.from_pretrained(
                args.model_name_or_path,
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True,
                device_map=args.device,
                trust_remote_code=True,
            )
        processor = AutoTokenizer.from_pretrained(
            args.model_name_or_path, trust_remote_code=True
        )
    model_unmerged = PeftModel.from_pretrained(
        model,
        args.adapter_name_or_path,
        device_map=args.device,
        is_trainable=False,
        trust_remote_code=True,
    )
    model_merged = model_unmerged.merge_and_unload()
    print(f"Saving to {args.save_dir}...")
    model_merged.save_pretrained(
        args.save_dir,
        safe_serialization=args.safe_serialization,
        max_shard_size=args.max_shard_size,
    )
    processor.save_pretrained(args.save_dir)
    print("All done!")


if __name__ == "__main__":
    main()
